from geesibling.core.types import Graph, Node
from geesibling.core.lib._graph import DataType,DeviceType,Device,search_policy
from geesibling.tools import log
from geesibling.adapters.pytorch.getTorchGraph import getTorchGraph
from geesibling.adapters.pytorch.graphToGeesiGraph import graphToGeeGraph
from geesibling.adapters.pytorch.opConvert import opNameConvert,opsToLayer
from geesibling.adapters.pytorch.geesiDeviceConfig import device_num
import torch

__doc__ = """
GeeSibling execute policy
Author: jirongting
datetime: 2024.05.17
version: 1 2024.05.17 first commit
"""

"""
add inputs for each geesibling node

Parameters:
 model_name - the name of model
 model - huggingface transformers model
 tokenizer - huggingface transformers tokenizer
 ngpu - the number of gpus
 policy - op set policy

Returns:
 hf_device_map - geesibling op set policy

"""
def GeeSiblingPolicy(model_name,model,tokenizer,ngpu,policy="sgp",pre = True,add_dim = None):

     # step1.1 get the input_shape
    text = "Replace me by any text you'd like."
    encoded_input = tokenizer(text, return_tensors='pt')
    input_size =encoded_input['input_ids']
    if not pre:
        input_size = input_size.unsqueeze(-1)
        new_shape = list(input_size.shape)
        new_shape[-1] = add_dim
        input_size = input_size.expand(new_shape)
        input_size = input_size.to(torch.float32)
    #input_size = input_size.to(torch.float32)
    # print(input_size.dtype)
    # print(input_size.shape)
    # step1.2 use symbolic_trace to trace the model and rewrite the interpreter to get shape of each node 
    traced = getTorchGraph(model,input_size,pre)
    #print(traced)
    # step1.3 graph convert (torch graph --> geesibling graph)
    geesiblingGraph = graphToGeeGraph(traced,model_name)
    #print(geesiblingGraph)
    # step2 search_policy(sgp)
    devices = device_num(ngpu=ngpu)
    geesi_device_map = search_policy(geesiblingGraph,devices,policy)
    # print(geesi_device_map)
    #print(geesi_device_map)
    # step3 opName convert
    hf_device_map = opNameConvert(geesi_device_map,model,model_name)
    hf_device_map = opsToLayer(hf_device_map)

    return hf_device_map
